# Copyright (c) HySoP 2011-2024
#
# This file is part of HySoP software.
# See "https://particle_methods.gricad-pages.univ-grenoble-alpes.fr/hysop-doc/"
# for further info.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from hysop.tools.io_utils import IOParams
from hysop.tools.htypes import check_instance, first_not_None
from hysop.tools.decorators import debug
from hysop.fields.continuous_field import Field
from hysop.topology.cartesian_descriptor import CartesianTopologyDescriptors
from hysop.operator.base.spectral_operator import (
SpectralOperatorBase,
EnergyPlotter,
EnergyDumper,
)
from hysop.symbolic.relational import Assignment
from hysop.symbolic.field import laplacian
[docs]
class PoissonOperatorBase(SpectralOperatorBase):
"""
Solves the poisson equation using a specific implementation.
"""
@debug
def __new__(
cls,
Fin,
Fout,
variables,
name=None,
pretty_name=None,
dump_energy=None,
dump_input_energy=None,
dump_output_energy=None,
plot_energy=None,
plot_input_energy=None,
plot_output_energy=None,
plot_inout_energy=None,
**kwds,
):
return super().__new__(
cls,
name=name,
pretty_name=pretty_name,
input_fields=None,
output_fields=None,
**kwds,
)
@debug
def __init__(
self,
Fin,
Fout,
variables,
name=None,
pretty_name=None,
dump_energy=None,
dump_input_energy=None,
dump_output_energy=None,
plot_energy=None,
plot_input_energy=None,
plot_output_energy=None,
plot_inout_energy=None,
**kwds,
):
"""
Initialize a n-dimensional Poisson operator base (using spectral methods).
Solves:
Laplacian(Fout) = Fin
Parameters
----------
Fout: Field
Input continuous field (rhs).
Fin: Field
Output continuous field (lhs), possibly inplace, same number of components as Fin.
variables: dict
Dictionary of fields as keys and topology descriptors as values.
dump_energy: IOParams, optional, defaults to None
Will set the default io parameter for all energy plotters.
dump_input_energy: IOParams, optional, defaults to None
Dump input field energy to a custom file. Defaults to no dump.
dump_output_energy: IOParams, optional, defaults to None
Dump output field energy to a custom file. Defaults to no dump.
plot_energy: IOParams, optional, defaults to None
Will set the default io parameter for all energy plotters.
plot_input_energy: IOParams, optional, defaults to None
Plot input field energy in a custom file. Defaults to no plot.
plot_output_energy: IOParams, optional, defaults to None
Plot output field energy in a custom file. Defaults to no plot.
plot_inout_energy: IOParams, optional, defaults to None
Plot input and output field energy on the same graph in a custom file.
Defaults to no plot.
kwds: dict, optional
Base class arguments.
"""
check_instance(Fin, Field)
check_instance(Fout, Field)
check_instance(variables, dict, keys=Field, values=CartesianTopologyDescriptors)
check_instance(dump_energy, (IOParams, int), allow_none=True)
check_instance(dump_input_energy, (IOParams, int), allow_none=True)
check_instance(dump_output_energy, (IOParams, int), allow_none=True)
check_instance(plot_energy, (IOParams, int), allow_none=True)
check_instance(plot_input_energy, (IOParams, int), allow_none=True)
check_instance(plot_output_energy, (IOParams, int), allow_none=True)
check_instance(plot_inout_energy, (IOParams, int), allow_none=True)
assert Fin.nb_components == Fout.nb_components
input_fields = {Fin: variables[Fin]}
output_fields = {Fout: variables[Fout]}
default_name = f"Poisson_{Fin.name}_{Fout.name}"
default_pretty_name = f"Poisson_{Fin.pretty_name}_{Fout.pretty_name}"
name = first_not_None(name, default_name)
pretty_name = first_not_None(name, default_pretty_name)
dump_input_E = first_not_None(dump_input_energy, dump_energy)
dump_output_E = first_not_None(dump_output_energy, dump_energy)
plot_input_E = first_not_None(plot_input_energy, plot_energy)
plot_output_E = first_not_None(plot_output_energy, plot_energy)
plot_inout_E = first_not_None(plot_inout_energy, plot_energy)
do_plot_inout_E = isinstance(plot_inout_E, IOParams) and (
plot_inout_E.frequency >= 0
)
_, compute_input_E_freqs = EnergyDumper.do_compute_energy(
dump_input_E, plot_input_E, plot_inout_E
)
_, compute_output_E_freqs = EnergyDumper.do_compute_energy(
dump_output_E, plot_output_E, plot_inout_E
)
super().__init__(
name=name,
pretty_name=pretty_name,
input_fields=input_fields,
output_fields=output_fields,
**kwds,
)
forward_transforms = ()
backward_transforms = ()
wave_numbers = ()
tg = self.new_transform_group()
for Fi, Fo in zip(Fin.fields, Fout.fields):
Ft = tg.require_forward_transform(
Fi,
custom_output_buffer="auto",
dump_energy=dump_input_E,
plot_energy=plot_input_E,
compute_energy_frequencies=compute_input_E_freqs,
)
Bt = tg.require_backward_transform(
Fo,
custom_input_buffer="auto",
matching_forward_transform=Ft,
dump_energy=dump_output_E,
plot_energy=plot_output_E,
compute_energy_frequencies=compute_output_E_freqs,
)
assert Ft.output_dtype == Bt.input_dtype
expr = Assignment(Bt.s, laplacian(Ft.s, Ft.s.frame))
kds = tg.push_expressions(expr)
forward_transforms += (Ft,)
backward_transforms += (Bt,)
wave_numbers += (kds,)
self.Fin = Fin
self.Fout = Fout
self.tg = tg
self.forward_transforms = forward_transforms
self.backward_transforms = backward_transforms
self.wave_numbers = wave_numbers
self.do_plot_inout_energy = do_plot_inout_E
self.plot_inout_energy_ioparams = plot_inout_E
self.input_energy_params = tuple(
Ft._energy_parameter for Ft in self.forward_transforms
)
self.output_energy_params = tuple(
Bt._energy_parameter for Bt in self.backward_transforms
)
[docs]
@debug
def discretize(self):
super().discretize()
self.dFin = self.get_input_discrete_field(self.Fin)
self.dFout = self.get_output_discrete_field(self.Fout)
all_dkds, all_nd_dkds = (), ()
for wn in self.wave_numbers:
dkds, nd_dkds = [
None,
] * len(wn), [
None,
] * len(wn)
for wi in wn:
idx, dkd, nd_dkd = self.tg.discrete_wave_numbers[wi]
dkds[idx] = dkd
nd_dkds[idx] = nd_dkd
all_dkds += (dkds,)
all_nd_dkds += (all_dkds,)
self.all_dkds = all_dkds
self.all_nd_dkds = all_nd_dkds
inout_energy_plotters = ()
if self.do_plot_inout_energy:
for dFin, dFout, Pin, Pout in zip(
self.dFin.dfields,
self.dFout.dfields,
self.input_energy_params,
self.output_energy_params,
):
fname = f"{dFin.name}_{dFout.name}"
iname = f"{type(self).__name__}.input.{dFin.pretty_name}"
oname = f"{type(self).__name__}.output.{dFout.pretty_name}"
plt = EnergyPlotter(
energy_parameters={iname: Pin, oname: Pout},
fname=fname,
io_params=self.plot_inout_energy_ioparams,
)
inout_energy_plotters += (plt,)
self.inout_energy_plotters = inout_energy_plotters
[docs]
def plot(self, simulation):
if not self.do_plot_inout_energy:
return
for plt in self.inout_energy_plotters:
plt.update(simulation=simulation, wait_for=None)